#include <fstream>
#include <string>
#include <numeric>
#include "transform.h"
#include "contextchain.h"

CompiledLinearTransform::CompiledLinearTransform(const std::shared_ptr<const SlotRing> &slotRing,
    std::vector<poly> coefficients)
    : slotRing(slotRing), coefficients(std::move(coefficients))
{
    assert(slotRing->G2Ord() == 2);
    assert(this->coefficients.size() == G1SubgroupOrder() || this->coefficients.size() == 2 * G1SubgroupOrder());
}

SlotRing::RawAuto CompiledLinearTransform::Automorphism(size_t index) const
{
    const size_t g1Index = index % G1SubgroupOrder();
    const size_t g2Index = (index % coefficients.size()) / G1SubgroupOrder();
    return slotRing->RawAuto_(g1Index * (slotRing->G1Ord() / G1SubgroupOrder()), g2Index);
}

SlotRing::RawAuto CompiledLinearTransform::DifferenceAutomorphism(size_t from, size_t to) const
{
    const size_t fromG1Index = (from % G1SubgroupOrder());
    const size_t fromG2Index = (from / G1SubgroupOrder()) % G2SubgroupOrder();
    const size_t toG1Index = (to % G1SubgroupOrder());
    const size_t toG2Index = (to / G1SubgroupOrder()) % G2SubgroupOrder();
    size_t diffG1Index = toG1Index + G1SubgroupOrder() - fromG1Index;
    if (diffG1Index >= G1SubgroupOrder()) {
        diffG1Index -= G1SubgroupOrder();
    }
    size_t diffG2Index = toG2Index + G2SubgroupOrder() - fromG2Index;
    if (diffG2Index >= G2SubgroupOrder()) {
        diffG2Index -= G2SubgroupOrder();
    }
    return Automorphism(diffG1Index + diffG2Index * G1SubgroupOrder());
}

SlotRing::RawAuto CompiledLinearTransform::ReverseAutomorphism(size_t index) const
{
    return DifferenceAutomorphism(index, 0);
}

uint64_t CompiledLinearTransform::G1SubgroupOrder() const
{
    return slotRing->G1Ord();
}

uint64_t CompiledLinearTransform::G2SubgroupOrder() const
{
    if (coefficients.size() == G1SubgroupOrder()) {
        return 1;
    } else {
        assert(coefficients.size() == 2 * G1SubgroupOrder());
        return 2;
    }
}

void CompiledLinearTransform::AddScaledTransform(const poly &scaling, const SlotRing::Rotation &rotation,
    const SlotRing::Frobenius &frobenius)
{
    {
        const uint64_t g1Power =
            (std::get<0>(rotation.GetForwardG1G2Decomp()) + std::get<0>(frobenius.GetG1G2Decomp())) % slotRing->G1Ord();
        const uint64_t g2Power =
            (std::get<1>(rotation.GetForwardG1G2Decomp()) + std::get<1>(frobenius.GetG1G2Decomp())) % slotRing->G2Ord();
        assert((g1Power * G1SubgroupOrder()) % slotRing->G1Ord() == 0);
        const size_t index = g1Power * G1SubgroupOrder() / slotRing->G1Ord() + g2Power * G1SubgroupOrder();
        poly mask = slotRing->RawAuto_(g1Power, g2Power)(rotation.GetForwardMask());
        poly add = PolyMulMod(scaling, mask, slotRing->R().scalarMod, slotRing->R().polyMod);

        PolyAdd(coefficients[index], add, slotRing->R().scalarMod);
    }
    {
        const uint64_t g1Power =
            (std::get<0>(rotation.GetBackwardG1G2Decomp()) + std::get<0>(frobenius.GetG1G2Decomp())) %
            slotRing->G1Ord();
        const uint64_t g2Power =
            (std::get<1>(rotation.GetBackwardG1G2Decomp()) + std::get<1>(frobenius.GetG1G2Decomp())) %
            slotRing->G2Ord();
        assert((g1Power * G1SubgroupOrder()) % slotRing->G1Ord() == 0);
        const size_t index = g1Power * G1SubgroupOrder() / slotRing->G1Ord() + g2Power * G1SubgroupOrder();
        poly mask = slotRing->RawAuto_(g1Power, g2Power)(rotation.GetBackwardMask());
        poly add = PolyMulMod(scaling, mask, slotRing->R().scalarMod, slotRing->R().polyMod);

        PolyAdd(coefficients[index], add, slotRing->R().scalarMod);
    }
}

size_t CompiledLinearTransform::BabystepAutomorphismCount() const
{
    const size_t automorphismCountLog2 = Log2Exact(coefficients.size());
    return static_cast<size_t>(1) << (automorphismCountLog2 - automorphismCountLog2 / 2);
}

size_t CompiledLinearTransform::GiantstepAutomorphismCount() const
{
    const size_t automorphismCountLog2 = Log2Exact(coefficients.size());
    return static_cast<size_t>(1) << (automorphismCountLog2 / 2);
}

void CompiledLinearTransform::FixCoefficientShift()
{
    for (size_t i = 0; i < GiantstepAutomorphismCount(); ++i) {
        for (size_t j = 0; j < BabystepAutomorphismCount(); ++j) {
            const auto automorphism = ReverseAutomorphism(i * BabystepAutomorphismCount());
            this->coefficients[i * BabystepAutomorphismCount() + j] =
                automorphism(this->coefficients[i * BabystepAutomorphismCount() + j]);
        }
    }
}

std::vector<poly> CompiledLinearTransform::CompileFrobenius(
    const std::unordered_map<std::tuple<size_t, size_t>, uint64_t> &sparseTransformMatrix,
    const SlotRing::SubringView &ring, uint64_t p, size_t N, const NegacyclicPowerTable &powerTable)
{
    size_t d = ring.polyMod.n;
    poly z = { 1 };

    const poly &xPowerN = ring.polyMod.xPowerN;
    const uint64_t &modT = ring.scalarMod;
    assert(SlotRing::IsPolySparse(xPowerN, d / 2));

    uint64_t a = 0;
    if (d / 2 < xPowerN.size()) {
        a = xPowerN[d / 2];
    }
    const uint64_t b = xPowerN[0];

    // trace of x^0 = 1
    const uint64_t u0 = d % modT;
    // trace of x^(d/2)
    const uint64_t u1 = MultiplyUintMod(a, d / 2, modT);
    // trace of x^d
    const uint64_t u2 =
        AddUintMod(MultiplyUintMod(MultiplyUintMod(a, a, modT), d / 2, modT), MultiplyUintMod(d, b, modT), modT);
    // trace of x^(3d/2)
    const uint64_t u3 = MultiplyUintMod(
        AddUintMod(ModExp(a, 3, modT), MultiplyUintMod(3, MultiplyUintMod(a, b, modT), modT), modT), d / 2, modT);
    const std::array<uint64_t, 4> us = { u0, u1, u2, u3 };
    // factor to scale with at the end
    uint64_t globalFactor = NegateUintMod(
        InvMod_(SubUintMod(MultiplyUintMod(u1, u3, modT), MultiplyUintMod(u2, u2, modT), modT), modT), modT);

    uint64_t Mod2N(2 * N);
    poly tmp;
    std::vector<poly> results;
    for (size_t m = 0; m < d; ++m) {
        poly result;
        for (const auto &entry : sparseTransformMatrix) {
            size_t k = std::get<0>(entry.first);
            size_t l = std::get<1>(entry.first);
            assert(k <= d);
            assert(l <= d);
            int64_t delta = (2 * l) / d + 1;
            int64_t pm = ModExp(p % Mod2N, m, Mod2N); // reduce
            tmp.clear();
            PolyAdd(tmp, powerTable[static_cast<int64_t>(k) + delta * pm * d / 2 - l * pm], modT, us[4 - delta]);
            PolySub(tmp, powerTable[static_cast<int64_t>(k) + (delta + 1) * pm * d / 2 - l * pm], modT, us[3 - delta]);
            uint64_t factor = entry.second % modT;
            if (delta % 2 != 0) {
                factor = NegateUintMod(factor, modT);
            }
            PolyAdd(result, tmp, modT, factor);
        }
        PolyScale(result, globalFactor, modT);
        results.push_back(std::move(result));
    }

    return results;
}

CompiledLinearTransform CompiledLinearTransform::ScalarSlotsToFirstCoefficients(
    std::shared_ptr<const SlotRing> slotRing)
{
    const NegacyclicPowerTable slotPowertable(slotRing->slot(), slotRing->slot().Generator(), slotRing->N());
    auto basisTransformMatrix = [&slotRing, &slotPowertable](
        std::unordered_map<std::tuple<size_t, size_t>, uint64_t> &block, size_t blockRow, size_t blockCol) {
        const size_t powerOfX = blockCol;
        const size_t outSlot = blockRow;
        const size_t powerOfZeta = MultiplyUintMod(powerOfX,
            InvMod_(std::get<0>(slotRing->Rotate(outSlot).GaloisElements()), slotRing->IndexMod()),
            slotRing->IndexMod());
        // this entry in the slot outSlot corresponds to the coset of x^powerOfX
        const poly entry = slotPowertable[powerOfZeta];
        for (size_t i = 0; i < slotRing->SlotRank(); ++i) {
            if (entry[i] != 0) {
                block[std::make_tuple(i, 0)] = entry[i];
            }
        }
    };
    return CompiledLinearTransform::CompileSlotBasis(slotRing, basisTransformMatrix);
}

CompiledLinearTransform CompiledLinearTransform::FirstCoefficientsToScalarSlots(
    std::shared_ptr<const SlotRing> slotRing)
{
    auto basisTransformMatrix = [&slotRing](std::unordered_map<std::tuple<size_t, size_t>, uint64_t> &block,
        size_t blockRow, size_t blockCol) {
        for (size_t k = 0; k < slotRing->SlotRank(); ++k) {
            const size_t power = SubUintMod(blockRow,
                MultiplyUintMod(std::get<0>(slotRing->Rotate(blockCol).GaloisElements()), k, slotRing->IndexMod()),
                slotRing->IndexMod());
            uint64_t value;
            if (power >= slotRing->N()) {
                value = NegateUintMod(slotRing->SlotOne(blockCol)[power - slotRing->N()], slotRing->R().scalarMod);
            } else {
                value = slotRing->SlotOne(blockCol)[power];
            }
            if (value != 0) {
                block[std::make_tuple(0, k)] = value;
            }
        }
    };
    return CompiledLinearTransform::CompileSlotBasis(slotRing, basisTransformMatrix);
}

CompiledSubringLinearTransform CompiledLinearTransform::InRing() &&
{
    return CompiledSubringLinearTransform(std::move(*this), slotRing);
}

NegacyclicPowerTable::NegacyclicPowerTable(const SlotRing::SubringView &ring, poly generator,
    uint64_t halfOrderGenerator)
    : ring(ring), generator(), N(halfOrderGenerator)
{
    // the "identity" does not have to be 1 - it can be 1 in the current slot, and 0 in the others
    poly current = PolyExponentiateMod(generator, 2 * halfOrderGenerator, ring.scalarMod, ring.polyMod);
    content.reserve(N);
    for (size_t i = 0; i < N; ++i) {
        poly next = PolyMulMod(current, generator, ring.scalarMod, ring.polyMod);
        content.push_back(std::move(current));
        current = std::move(next);
    }
    this->generator = std::move(generator);
}

poly NegacyclicPowerTable::operator[](int64_t i) const
{
    int64_t index = i % static_cast<int64_t>(2 * N);
    if (index < 0) {
        index += 2 * N;
    }
    if (index < static_cast<int64_t>(N)) {
        return content[index];
    } else {
        poly result = content[index - N];
        PolyScale(result, NegateUintMod(1, ring.scalarMod), ring.scalarMod);
        return result;
    }
}

void TestFirstCoeffsToScalarSlots()
{
    std::shared_ptr<SlotRing> slotRing = std::make_shared<SlotRing>(P257TestParameters()->PowerXSubring(7));
    std::shared_ptr<SlotRing> smallerSlotRing = std::make_shared<SlotRing>(slotRing->PowerXSubring(1));
    CompiledLinearTransform t = CompiledLinearTransform::FirstCoefficientsToScalarSlots(smallerSlotRing);
    CompiledSubringLinearTransform transform(std::move(t), slotRing);

    poly a = { 1, 0, 2, 0, 4, 0 };
    a.resize(slotRing->N());
    poly b = transform(a);
    assert(slotRing->ExtractSlotValue(b, 0)[0] == 1);
    assert(slotRing->ExtractSlotValue(b, 1)[0] == 2);
    assert(slotRing->ExtractSlotValue(b, 2)[0] == 4);
    for (size_t i = 3; i < slotRing->SlotGroupLen(); ++i) {
        assert(slotRing->ExtractSlotValue(b, i)[0] == 0);
    }
    std::cout << "TestFirstCoeffsToScalarSlots(): success" << std::endl;
}

void TestCompileSlotBasis()
{
    std::shared_ptr<SlotRing> slotRing = SmallTestParameters();
    auto matrix = [&slotRing](std::unordered_map<std::tuple<size_t, size_t>, uint64_t> &slotwiseMatrix, size_t i,
        size_t j) {
        if (i == j) {
            for (size_t l = 0; l < slotRing->SlotRank(); ++l) {
                slotwiseMatrix[std::make_tuple(l, l)] = 1;
            }
        } else if (i == j + 16 && i < 32) {
            slotwiseMatrix[std::make_tuple(0, 0)] = 1;
        } else if (j == i + 16 && j < 32) {
            slotwiseMatrix[std::make_tuple(1, 1)] = 1;
        }
    };
    CompiledLinearTransform t = CompiledLinearTransform::CompileSlotBasis(slotRing, matrix);
    CompiledSubringLinearTransform transform = std::move(t).InRing();

    poly a = slotRing->FromSlotValue({ 1, 1 }, 0);
    PolyAdd(a, slotRing->FromSlotValue({ 2, 0, 3 }, 1), slotRing->R().scalarMod);
    PolyAdd(a, slotRing->FromSlotValue({ 4, 0, 1 }, 16), slotRing->R().scalarMod);
    PolyAdd(a, slotRing->FromSlotValue({ 0, 1 }, 17), slotRing->R().scalarMod);

    poly resultPoly = transform(a);

    poly expected;
    PolyAdd(expected, slotRing->FromSlotValue({ 1, 1 }, 0), slotRing->R().scalarMod);
    PolyAdd(expected, slotRing->FromSlotValue({ 2, 1, 3 }, 1), slotRing->R().scalarMod);
    PolyAdd(expected, slotRing->FromSlotValue({ 5, 0, 1 }, 16), slotRing->R().scalarMod);
    PolyAdd(expected, slotRing->FromSlotValue({ 2, 1 }, 17), slotRing->R().scalarMod);

    assert(resultPoly == expected);
    std::cout << "TestCompileSlotBasis(): success" << std::endl;
}

void TestApplyCiphertext()
{
    std::shared_ptr<SlotRing> slotRing = SmallTestParameters();
    auto matrix = [&slotRing](std::unordered_map<std::tuple<size_t, size_t>, uint64_t> &slotwiseMatrix, size_t i,
        size_t j) {
        if (i == j) {
            for (size_t l = 0; l < slotRing->SlotRank(); ++l) {
                slotwiseMatrix[std::make_tuple(l, l)] = 1;
            }
        } else if (i == j + 16 && i < 32) {
            slotwiseMatrix[std::make_tuple(0, 0)] = 1;
        } else if (j == i + 16 && j < 32) {
            slotwiseMatrix[std::make_tuple(1, 1)] = 1;
        }
    };
    CompiledLinearTransform t = CompiledLinearTransform::CompileSlotBasis(slotRing, matrix);
    CompiledSubringLinearTransform transform = std::move(t).InRing();

    Context context(9, { 1099511590913 }, { 1099511592961 }, slotRing->Prime(), 3.2);

    KeyGenerator keygen(&context);
    std::pair<SecretKey *, PublicKey *> tmp = keygen.GenKeyPair();
    SecretKey sk(*tmp.first);
    PublicKey pk(*tmp.second);
    vector<uint64_t> galoisVector;
    for (int i = 0; i < transform.GaloisElements().size(); i++) {
        galoisVector.push_back(transform.GaloisElements()[i]);
    }
    GaloisKeys gk(*keygen.GenGaloisKeys(galoisVector, &sk));

    PkEncryptor encryptor(&context, &pk);
    Evaluator evaluator(&context);
    Decryptor decryptor(&context, &sk);

    poly a = slotRing->FromSlotValue({ 1, 1 }, 0);
    PolyAdd(a, slotRing->FromSlotValue({ 2, 0, 3 }, 1), slotRing->R().scalarMod);
    PolyAdd(a, slotRing->FromSlotValue({ 4, 0, 1 }, 16), slotRing->R().scalarMod);
    PolyAdd(a, slotRing->FromSlotValue({ 0, 1 }, 17), slotRing->R().scalarMod);
    Plaintext xPlain(a, context.MaxLevel());
    evaluator.lightEncoder->ScaleUp(&xPlain, &xPlain);

    Ciphertext xEnc(&context, 1, context.MaxLevel());
    encryptor.Encrypt(&xPlain, &xEnc);

    Ciphertext resultEnc(&context, 1, context.MaxLevel());
    transform.ApplyCiphertext(xEnc, evaluator, gk, resultEnc);

    Plaintext result(&context, context.MaxLevel());
    decryptor.Decrypt(&resultEnc, &result);

    evaluator.lightEncoder->ScaleDown(&result, &result);

    poly resultPoly(&result.value->Coeffs[0][0], &result.value->Coeffs[0][0] + result.CoeffCount());
    resultPoly.resize(slotRing->N());
    poly expected;
    PolyAdd(expected, slotRing->FromSlotValue({ 1, 1 }, 0), slotRing->R().scalarMod);
    PolyAdd(expected, slotRing->FromSlotValue({ 2, 1, 3 }, 1), slotRing->R().scalarMod);
    PolyAdd(expected, slotRing->FromSlotValue({ 5, 0, 1 }, 16), slotRing->R().scalarMod);
    PolyAdd(expected, slotRing->FromSlotValue({ 2, 1 }, 17), slotRing->R().scalarMod);

    assert(resultPoly == expected);
    std::cout << "TestApplyCiphertext(): success" << std::endl;
}

void TestApplyCiphertextSubring()
{
    std::shared_ptr<SlotRing> largeSlotRing = SmallTestParameters();
    std::shared_ptr<SlotRing> slotRing = std::make_shared<SlotRing>(largeSlotRing->PowerXSubring(2));
    auto matrix = [&slotRing](std::unordered_map<std::tuple<size_t, size_t>, uint64_t> &slotwiseMatrix, size_t i,
        size_t j) {
        if (i == j) {
            for (size_t l = 0; l < 2; ++l) {
                slotwiseMatrix[std::make_tuple(l, l)] = 1;
            }
        } else if (i == j + 16) {
            slotwiseMatrix[std::make_tuple(0, 0)] = 1;
        } else if (j == i + 16) {
            slotwiseMatrix[std::make_tuple(1, 1)] = 1;
        }
    };
    CompiledLinearTransform t = CompiledLinearTransform::CompileSlotBasis(slotRing, matrix);
    CompiledSubringLinearTransform transform(std::move(t), largeSlotRing);

    Context context(9, { 1099511590913 }, { 1099511592961 }, largeSlotRing->Prime(), 3.2);
    KeyGenerator keygen(&context);
    std::pair<SecretKey *, PublicKey *> tmp = keygen.GenKeyPair();
    SecretKey sk(*tmp.first);
    PublicKey pk(*tmp.second);
    vector<uint64_t> galoisVector;
    for (int i = 0; i < transform.GaloisElements().size(); i++) {
        galoisVector.push_back(transform.GaloisElements()[i]);
    }
    GaloisKeys gk(*keygen.GenGaloisKeys(galoisVector, &sk));

    PkEncryptor encryptor(&context, &pk);
    Evaluator evaluator(&context);
    Decryptor decryptor(&context, &sk);

    poly a = largeSlotRing->FromSlotValue({ 1, 0, 0, 0, 1 }, 0);
    PolyAdd(a, largeSlotRing->FromSlotValue({ 2 }, 1), largeSlotRing->R().scalarMod);
    PolyAdd(a, largeSlotRing->FromSlotValue({ 4 }, 16), largeSlotRing->R().scalarMod);
    PolyAdd(a, largeSlotRing->FromSlotValue({ 0, 0, 0, 0, 1 }, 17), largeSlotRing->R().scalarMod);
    Plaintext xPlain(a, context.MaxLevel());

    evaluator.lightEncoder->ScaleUp(&xPlain, &xPlain);

    Ciphertext xEnc(&context, 1, context.MaxLevel());
    encryptor.Encrypt(&xPlain, &xEnc);

    Ciphertext resultEnc(&context, 1, context.MaxLevel());
    transform.ApplyCiphertext(xEnc, evaluator, gk, resultEnc);

    Plaintext result(&context, context.MaxLevel());
    decryptor.Decrypt(&resultEnc, &result);

    evaluator.lightEncoder->ScaleDown(&result, &result);

    poly resultPoly(&result.value->Coeffs[0][0], &result.value->Coeffs[0][0] + result.CoeffCount());
    resultPoly.resize(largeSlotRing->N());
    poly expected;
    PolyAdd(expected, largeSlotRing->FromSlotValue({ 1, 0, 0, 0, 1 }, 0), largeSlotRing->R().scalarMod);
    PolyAdd(expected, largeSlotRing->FromSlotValue({ 2, 0, 0, 0, 1 }, 1), largeSlotRing->R().scalarMod);
    PolyAdd(expected, largeSlotRing->FromSlotValue({ 5 }, 16), largeSlotRing->R().scalarMod);
    PolyAdd(expected, largeSlotRing->FromSlotValue({ 2, 0, 0, 0, 1 }, 17), largeSlotRing->R().scalarMod);

    assert(resultPoly == expected);
    std::cout << "TestApplyCiphertextSubring(): success" << std::endl;
}

SlotRing::RawAuto CompiledSubringLinearTransform::Automorphism(size_t index) const
{
    assert(slotRing->SlotGroupLen() == subringTransform.slotRing->SlotGroupLen());
    const size_t g1Index = index % subringTransform.G1SubgroupOrder();
    const size_t g2Index = (index % subringTransform.coefficients.size()) / subringTransform.G1SubgroupOrder();
    return slotRing->RawAuto_(g1Index * (subringTransform.slotRing->G1Ord() / subringTransform.G1SubgroupOrder()),
        g2Index);
}

SlotRing::RawAuto CompiledSubringLinearTransform::DifferenceAutomorphism(size_t from, size_t to) const
{
    const size_t fromG1Index = (from % subringTransform.G1SubgroupOrder());
    const size_t fromG2Index = (from % subringTransform.coefficients.size()) / subringTransform.G1SubgroupOrder();
    const size_t toG1Index = (to % subringTransform.G1SubgroupOrder());
    const size_t toG2Index = (to % subringTransform.coefficients.size()) / subringTransform.G1SubgroupOrder();
    size_t diffG1Index = toG1Index + subringTransform.G1SubgroupOrder() - fromG1Index;
    if (diffG1Index >= subringTransform.G1SubgroupOrder()) {
        diffG1Index -= subringTransform.G1SubgroupOrder();
    }
    size_t diffG2Index = toG2Index + subringTransform.G2SubgroupOrder() - fromG2Index;
    if (diffG2Index >= subringTransform.G2SubgroupOrder()) {
        diffG2Index -= subringTransform.G2SubgroupOrder();
    }
    return Automorphism(diffG1Index + diffG2Index * subringTransform.G1SubgroupOrder());
}

size_t CompiledSubringLinearTransform::BabystepAutomorphismCount() const
{
    return subringTransform.BabystepAutomorphismCount();
}

size_t CompiledSubringLinearTransform::GiantstepAutomorphismCount() const
{
    return subringTransform.GiantstepAutomorphismCount();
}

CompiledSubringLinearTransform::CompiledSubringLinearTransform(CompiledLinearTransform &&transform,
    std::shared_ptr<const SlotRing> newRing)
    : slotRing(newRing), subringTransform(std::move(transform))
{}

CompiledSubringLinearTransform CompiledSubringLinearTransform::SlotsToCoeffs(std::shared_ptr<const SlotRing> slotRing)
{
    assert(slotRing->IsDHalfSparse());
    if (slotRing->IsDSparse()) {
        std::shared_ptr<SlotRing> reducedSlotRing =
            std::make_shared<SlotRing>(slotRing->PowerXSubring(Log2Exact(slotRing->SlotRank())));
        CompiledLinearTransform t = CompiledLinearTransform::ScalarSlotsToFirstCoefficients(reducedSlotRing);
        return CompiledSubringLinearTransform(std::move(t), slotRing);
    } else {
        std::shared_ptr<SlotRing> reducedSlotRing =
            std::make_shared<SlotRing>(slotRing->PowerXSubring(Log2Exact(slotRing->SlotRank() / 2)));
        CompiledLinearTransform t = CompiledLinearTransform::ScalarSlotsToFirstCoefficients(reducedSlotRing);
        return CompiledSubringLinearTransform(std::move(t), slotRing);
    }
}

CompiledSubringLinearTransform CompiledSubringLinearTransform::CoeffsToSlots(std::shared_ptr<const SlotRing> slotRing)
{
    assert(slotRing->IsDHalfSparse());
    if (slotRing->IsDSparse()) {
        std::shared_ptr<SlotRing> reducedSlotRing =
            std::make_shared<SlotRing>(slotRing->PowerXSubring(Log2Exact(slotRing->SlotRank())));
        CompiledLinearTransform t = CompiledLinearTransform::FirstCoefficientsToScalarSlots(reducedSlotRing);
        return CompiledSubringLinearTransform(std::move(t), slotRing);
    } else {
        std::shared_ptr<SlotRing> reducedSlotRing =
            std::make_shared<SlotRing>(slotRing->PowerXSubring(Log2Exact(slotRing->SlotRank() / 2)));
        CompiledLinearTransform t = CompiledLinearTransform::FirstCoefficientsToScalarSlots(reducedSlotRing);
        return CompiledSubringLinearTransform(std::move(t), slotRing);
    }
}

poly CompiledSubringLinearTransform::operator () (const poly &x) const
{
    std::vector<poly> precomputedValues;
    precomputedValues.emplace_back(x);
    for (size_t i = 1; i < BabystepAutomorphismCount(); ++i) {
        // 开始移除低位, 这样具有重新旋转和计算 Frobenius_ 的效果，
        size_t baseElementIndex = i - (static_cast<size_t>(1) << HighestDividingPower2(i));
        assert(baseElementIndex < i);
        SlotRing::RawAuto automorphismToApply = DifferenceAutomorphism(baseElementIndex, i);
        precomputedValues.push_back(automorphismToApply(precomputedValues[baseElementIndex]));
    }

    poly result;
    for (size_t i = 0; i < GiantstepAutomorphismCount(); ++i) {
        poly current;
        for (size_t j = 0; j < BabystepAutomorphismCount(); ++j) {
            if (!IsZero(subringTransform.coefficients[i * BabystepAutomorphismCount() + j])) {
                poly coeff;
                slotRing->FromPowerXSubring(*subringTransform.slotRing,
                    subringTransform.coefficients[i * BabystepAutomorphismCount() + j], coeff);
                poly addition = PolyMulMod(coeff, precomputedValues[j], slotRing->R().scalarMod, slotRing->R().polyMod);
                PolyAdd(current, addition, slotRing->R().scalarMod);
            }
        }
        SlotRing::RawAuto automorphismToApply = Automorphism(i * BabystepAutomorphismCount());
        PolyAdd(result, automorphismToApply(current), slotRing->R().scalarMod);
    }
    return result;
}

/*
 * @brief 线性变换主函数，使用BSGS算法执行密文乘以一个DFT矩阵。
 *
 * @param in 输入的密文
 * @param eval 同态计算器，包含自同构，密文乘等函数
 * @param gk galois密钥
 * @param result 线性变换后的结果
 */
void CompiledSubringLinearTransform::ApplyCiphertext(Ciphertext &in, Evaluator &eval, GaloisKeys &gk,
    Ciphertext &result) const
{
    if (coefficientsPlain.size() == 0) {
        poly coeff;
        for (const auto &c : subringTransform.coefficients) {
            slotRing->FromPowerXSubring(*subringTransform.slotRing, c, coeff);
            coefficientsPlain.emplace_back(coeff, eval.params->MaxLevel());
        }
    } else {
        assert(coefficientsPlain.size() == subringTransform.coefficients.size());
    }

    std::vector<Ciphertext> precomputedValues;
    precomputedValues.emplace_back(in);
    Ciphertext tmp(eval.params->N(), 1, eval.params->MaxLevel());
    for (size_t i = 1; i < BabystepAutomorphismCount(); ++i) {
        Ciphertext tmp1(eval.params->N(), 1, eval.params->MaxLevel());
        // 开始移除低位, 这样具有重新旋转和计算 Frobenius_ 的效果，
        size_t baseElementIndex = i - (static_cast<size_t>(1) << HighestDividingPower2(i));
        assert(baseElementIndex < i);
        SlotRing::RawAuto automorphismToApply = DifferenceAutomorphism(baseElementIndex, i);
        automorphismToApply.ApplyCiphertext(precomputedValues[baseElementIndex], eval, gk, tmp1);
        precomputedValues.push_back(tmp1);
    }
    bool isResultSet = false;
    for (size_t i = 0; i < GiantstepAutomorphismCount(); ++i) {
        Ciphertext current(eval.params->N(), 1, eval.params->MaxLevel());
        bool is_current_set = false;
        for (size_t j = 0; j < BabystepAutomorphismCount(); ++j) {
            if (!IsZero(subringTransform.coefficients[i * BabystepAutomorphismCount() + j])) {
                if (is_current_set) {
                    eval.Mul(&precomputedValues[j], &coefficientsPlain[i * BabystepAutomorphismCount() + j], &tmp,
                        MultiType::plaintextRingT);
                    eval.Add(&current, &tmp, &current);
                } else {
                    eval.Mul(&precomputedValues[j], &coefficientsPlain[i * BabystepAutomorphismCount() + j], &current,
                        MultiType::plaintextRingT);
                    is_current_set = true;
                }
            }
        }
        if (is_current_set) {
            SlotRing::RawAuto automorphismToApply = Automorphism(i * BabystepAutomorphismCount());
            automorphismToApply.ApplyCiphertext(current, eval, gk, tmp);
            if (!isResultSet) {
                result = tmp;
                isResultSet = true;
            } else {
                eval.Add(&result, &tmp, &result);
            }
        }
    }
}

std::vector<uint32_t> CompiledSubringLinearTransform::GaloisElements() const
{
    std::vector<uint32_t> result;
    for (size_t i = 1; i < BabystepAutomorphismCount(); ++i) {
        // 开始移除低位, 这样具有重新旋转和计算 Frobenius_ 的效果，
        size_t baseElementIndex = i - (static_cast<size_t>(1) << HighestDividingPower2(i));
        assert(baseElementIndex < i);
        SlotRing::RawAuto automorphismToApply = DifferenceAutomorphism(baseElementIndex, i);
        result.push_back(automorphismToApply.galois_element());
    }
    for (size_t i = 1; i < GiantstepAutomorphismCount(); ++i) {
        const auto automorphism = this->Automorphism(i * BabystepAutomorphismCount());
        result.push_back(automorphism.galois_element());
    }
    return result;
}

CompiledLinearTransform &&CompiledSubringLinearTransform::transform() &&
{
    return std::move(this->subringTransform);
}
